from __future__ import annotations

import os
from pathlib import Path
from typing import List

import fitz  # PyMuPDF
from PIL import Image

# Try to silence MuPDF error prints
try:
	fitz.TOOLS.mupdf_display_errors(False)
except Exception:
	pass

try:
	import pypdfium2 as pdfium
	_HAS_PDFIUM = True
except Exception:
	_HAS_PDFIUM = False


def _open_pdf_with_fallback(pdf_path: str) -> fitz.Document:
	"""Open a PDF, retrying with stream-based open if xref is damaged."""
	try:
		return fitz.open(pdf_path)
	except Exception:
		with open(pdf_path, "rb") as f:
			data = f.read()
		return fitz.open(stream=data, filetype="pdf")


def _render_with_pymupdf(pdf_path: str, output_dir: str, dpi: int) -> List[str]:
	png_paths: List[str] = []
	doc = _open_pdf_with_fallback(pdf_path)
	with doc:
		for page_index in range(doc.page_count):
			page = doc.load_page(page_index)
			pix = page.get_pixmap(dpi=dpi, alpha=False)
			out_path = os.path.join(output_dir, f"page_{page_index + 1:03d}.png")
			pix.save(out_path)
			png_paths.append(out_path)
	return png_paths


def _render_with_pdfium(pdf_path: str, output_dir: str, dpi: int) -> List[str]:
	png_paths: List[str] = []
	pdf = pdfium.PdfDocument(pdf_path)
	scale = dpi / 72.0
	for i in range(len(pdf)):
		page = pdf[i]
		bitmap = page.render(scale=scale, rotation=0)
		pil_image = bitmap.to_pil()
		out_path = os.path.join(output_dir, f"page_{i + 1:03d}.png")
		pil_image.save(out_path, format="PNG")
		png_paths.append(out_path)
	return png_paths


def convert_pdf_to_pngs(pdf_path: str, output_dir: str, dpi: int = 200) -> List[str]:
	"""Convert each PDF page to a PNG file.

	Args:
		pdf_path: Path to the input PDF.
		output_dir: Directory to write PNG files.
		dpi: Target DPI for rasterization.

	Returns:
		List of file paths to the generated PNG files in order.
	"""
	Path(output_dir).mkdir(parents=True, exist_ok=True)

	# Prefer pypdfium2 if available to avoid MuPDF xref errors
	if _HAS_PDFIUM:
		png_paths = _render_with_pdfium(pdf_path, output_dir, dpi)
	else:
		png_paths = _render_with_pymupdf(pdf_path, output_dir, dpi)

	# Validate generated images using Pillow
	valid_paths: List[str] = []
	for out_path in png_paths:
		try:
			with Image.open(out_path) as im:
				im.load()
			valid_paths.append(out_path)
		except Exception:
			try:
				os.remove(out_path)
			except Exception:
				pass

	return valid_paths
